/*
 * @Author: huangyupei huangyupei2021@ia.ac.cn
 * @Date: 2024-11-11 16:25:04
 * @LastEditors: huangyupei huangyupei2021@ia.ac.cn
 * @LastEditTime: 2025-03-19 23:00:38
 * @FilePath: /SelfCalib_OptiAcoustic/include/extrinsicRefinement.h
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */

#pragma once

#include "Thirdparty/g2o/g2o/core/base_vertex.h"
#include "Thirdparty/g2o/g2o/core/base_binary_edge.h"
#include "Thirdparty/g2o/g2o/core/block_solver.h"
#include "Thirdparty/g2o/g2o/core/optimization_algorithm_levenberg.h"
#include "Thirdparty/g2o/g2o/core/optimization_algorithm_gauss_newton.h"
#include "Thirdparty/g2o/g2o/core/optimization_algorithm_dogleg.h"
#include "Thirdparty/g2o/g2o/core/robust_kernel_impl.h"
#include "Thirdparty/g2o/g2o/solvers/linear_solver_dense.h"
#include "Thirdparty/g2o/g2o/solvers/linear_solver_eigen.h"
#include "Thirdparty/g2o/g2o/types/types_six_dof_expmap.h"
#include "Thirdparty/g2o/g2o/types/types_sba.h"

#include "Thirdparty/g2o/g2o/core/base_multi_edge.h"

#include <opencv2/opencv.hpp>
#include <opencv2/core/eigen.hpp>
#include "sophus/se3.hpp"

#include "readImageFeatures.h"

using namespace std;
using namespace Eigen;
using namespace Sophus;

// class VertexPose;
// class VertexPoint;

namespace ExtrinRefine {

// Sophus::SE3d 顶点类型，用于外参顶点、声呐位姿顶点
class VertexPose: public g2o::BaseVertex<6, Sophus::SE3d> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    // 重置(初始化)
    virtual void setToOriginImpl() override {   // 重写基类虚函数
        _estimate = Sophus::SE3d();
    }

    // 自定义更新方式：SE3李代数中的左乘
    virtual void oplusImpl(const double* update) override {
        Eigen::Matrix<double, 6, 1> update_eigen;
        update_eigen << update[0], update[1], update[2], update[3], update[4], update[5];
        _estimate = Sophus::SE3d::exp(update_eigen) * _estimate;
    }

    virtual bool read(istream &in) override {}

    virtual bool write(ostream &out) const override {}
};


// 三维点顶点类型, 用于相机恢复的三维点或者声呐BA恢复的三维点
class VertexPoint: public g2o::BaseVertex<3, Vector3d> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    VertexPoint() {}
    virtual void setToOriginImpl() override {
        _estimate = Vector3d(0, 0, 0);
    }

    virtual void oplusImpl(const double *update_) override {
        Eigen::Map<const Eigen::Vector3d> update(update_);
        _estimate += update;
    }

    virtual bool read(istream &in) {}
    virtual bool write(ostream &out) const {}
};

inline Vector2d CameraProjection(Vector3d P_ci_j) {
    Eigen::Matrix3d K;
    // double focal_length = 335;
    // K << focal_length, 0, 320,
    //            0, focal_length, 240,
    //            0, 0, 1;
    
    double focal_length = 1200;
    K << focal_length, 0, 960,
         0, focal_length, 540,
         0, 0, 1;

    Vector3d pos_pixel = K * P_ci_j;        // 3D点投影至像素平面
    pos_pixel /= pos_pixel[2];              // 转换为归一化坐标
    return pos_pixel.head<2>();
}


class EdgeCameraReproj: public g2o::BaseMultiEdge<2, Vector2d> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    // EdgeCameraReproj(const Vector3d& pos, const Matrix3d& K): _pos3d(pos), _K(K){}
    EdgeCameraReproj() {
        // 在这里指定顶点的数量，例如需要 3 个顶点
        resize(3);
    }

    // 计算边的误差
    virtual void computeError() override {
        auto v0 = (VertexPose *) _vertices[0];
        auto v1 = (VertexPose *) _vertices[1];
        auto v2 = (VertexPoint *) _vertices[2];
        Sophus::SE3d T_sc = v0->estimate();
        Sophus::SE3d T_si_s0 = v1->estimate();
        Vector3d P_c0_j = v2->estimate();

        Vector3d P_ci_j = T_sc.inverse() * T_si_s0 * T_sc * P_c0_j;
        Vector2d proj = CameraProjection(P_ci_j);
        _error = _measurement - proj;

        // cout << "_measurement: " << _measurement.transpose() << endl;
        // cout << "proj: " << proj.transpose() << endl;
    }

    // 使用数值求导，暂时不计算雅克比矩阵
    
    virtual bool read(istream &in) override {}

    virtual bool write(ostream &out) const override {}
};

Eigen::Vector2d SonarProjection(const Vector3d &point) {
    // ~ 声呐投影函数的实现，3D点为直角坐标形式
    double Xs = point[0], Ys = point[1], Zs = point[2];
    double r = sqrt(Xs*Xs + Ys*Ys + Zs*Zs), theta = atan2(Xs, Ys);
    return Vector2d(r * sin(theta), r*cos(theta));
}

class EdgeSonarReproj: public g2o::BaseMultiEdge<2, Vector2d> {
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    EdgeSonarReproj() {
        // 在这里指定顶点的数量，例如需要 3 个顶点
        resize(3);
    }

    virtual void computeError() override {
        auto v0 = (VertexPose *) _vertices[0];
        auto v1 = (VertexPose *) _vertices[1];
        auto v2 = (VertexPoint *) _vertices[2];
        Sophus::SE3d T_sc = v0->estimate();
        Sophus::SE3d T_ci_c0 = v1->estimate();
        Vector3d P_s0_j = v2->estimate();

        Vector3d P_si_j = T_sc * T_ci_c0 * T_sc.inverse() * P_s0_j;
        Vector2d proj = SonarProjection(P_si_j);
        _error = _measurement - proj;
    }

    // 使用数值求导，暂时不计算雅克比矩阵

    virtual bool read(istream &in) {}

    virtual bool write(ostream &out) const {}
}; 


void ExtrinsicRefineUseCam(Sophus::SE3d &T_sc0, vector<Vector3d> &vPointsCamera, vector<Sophus::SE3d> vSonarPoses, vector<vector<Vector2d>> vCameraMeasAll);

void ExtrinsicRefineUseSonar(Sophus::SE3d &T_sc0, vector<Vector3d> &vPointsSonar, vector<Sophus::SE3d> vCameraPoses, vector<vector<Vector2d>> vSonarMeasAll);

class EdgePoseAlignment : public g2o::BaseUnaryEdge<6, SE3d, g2o::VertexSE3Expmap> {
public:
    EdgePoseAlignment(const SE3d& cam_pose, const SE3d& sonar_pose)
        : camPose_(cam_pose), sonarPose_(sonar_pose) {}

    virtual void computeError() override {
        auto v = (g2o::VertexSE3Expmap *) _vertices[0];
        SE3d T_sc = SE3d(v->estimate().rotation(), v->estimate().translation());

        // 计算变换误差：T_s1_si * T_sc = T_sc * T_c1_ci
        SE3d error_se3 = sonarPose_ * T_sc * camPose_.inverse() * T_sc.inverse();
        // 李代数误差
        _error = error_se3.log();
    }

    bool read(std::istream&) override { return false; }
    bool write(std::ostream&) const override { return false; }

private:
    SE3d camPose_;
    SE3d sonarPose_;
};

// 轨迹一致性优化函数
void TrajConsistencyOptimization(const vector<Sophus::SE3d>& vCameraPoses, const vector<Sophus::SE3d>& vSonarPoses, Sophus::SE3d& T_sc);

}